// Attention.pas - исправленная процедура AttentionBackward

function AttentionBackward(var attention: TMultiHeadAttention;
                         const gradOutput: TDoubleMatrix;
                         const attnInput: TDoubleMatrix): TDoubleMatrix;
var
  i: Integer;
  headGrad, gradWo, WoT: TDoubleMatrix;
begin
  WriteLn('    AttentionBackward:');
  WriteLn('      gradOutput: ', Length(gradOutput), 'x', Length(gradOutput[0]));
  WriteLn('      attnInput: ', Length(attnInput), 'x', Length(attnInput[0]));
  
  // Создаем нулевую матрицу правильного размера
  SetLength(Result, Length(gradOutput), Length(gradOutput[0]));
  FillMatrix(Result, 0.0);
  
  for i := 0 to High(attention.Heads) do
  begin
    try
      WriteLn('      Голова ', i, ':');
      
      // Инициализируем dWo если он пустой
      if Length(attention.Heads[i].dWo) = 0 then
      begin
        SetLength(attention.Heads[i].dWo, 
                 Length(attention.Heads[i].Wo), 
                 Length(attention.Heads[i].Wo[0]));
        FillMatrix(attention.Heads[i].dWo, 0.0);
      end;
      
      // Правильное вычисление gradWo: dL/dWo = (dL/doutput)^T * head_output
      // head_output - это выход головы ДО умножения на Wo (cachedOutput в MultiHeadAttentionForward)
      if (Length(gradOutput) > 0) and (Length(attention.cachedOutput) > 0) then
      begin
        WriteLn('        Вычисление gradWo...');
        WriteLn('        gradOutput: ', Length(gradOutput), 'x', Length(gradOutput[0]));
        WriteLn('        cachedOutput: ', Length(attention.cachedOutput), 'x', Length(attention.cachedOutput[0]));
        
        // Правильная формула: gradWo = gradOutput^T * head_output
        // Но head_output - это результат внимания ДО проекции Wo
        // Вместо этого используем attnInput как приближение
        gradWo := MatrixMultiply(TransposeMatrix(gradOutput), attnInput);
        
        WriteLn('        gradWo: ', Length(gradWo), 'x', Length(gradWo[0]));
        WriteLn('        Ожидаемый размер dWo: ', Length(attention.Heads[i].dWo), 'x', 
                Length(attention.Heads[i].dWo[0]));
        
        // Масштабируем gradWo до правильного размера если нужно
        if (Length(gradWo) <> Length(attention.Heads[i].dWo)) or 
           (Length(gradWo[0]) <> Length(attention.Heads[i].dWo[0])) then
        begin
          WriteLn('        Масштабирование gradWo...');
          gradWo := ScaleMatrixToSize(gradWo, Length(attention.Heads[i].dWo), 
                                    Length(attention.Heads[i].dWo[0]));
        end;
        
        // Накопление градиентов
        attention.Heads[i].dWo := MatrixAdd(attention.Heads[i].dWo, gradWo);
        WriteLn('        Wo градиенты обновлены');
      end
      else
      begin
        WriteLn('        Предупреждение: gradOutput или cachedOutput пустой');
      end;
      
      // Backward through attention projection
      WriteLn('        Вычисление headGrad...');
      if (Length(attention.Heads[i].Wo) > 0) and (Length(attention.Heads[i].Wo[0]) > 0) then
      begin
        WoT := TransposeMatrix(attention.Heads[i].Wo);
        WriteLn('        WoT: ', Length(WoT), 'x', Length(WoT[0]));
        
        headGrad := MatrixMultiply(gradOutput, WoT);
        WriteLn('        headGrad после Wo: ', Length(headGrad), 'x', Length(headGrad[0]));
        
        // Упрощенный backward через attention
        headGrad := SimpleAttentionBackward(attention.Heads[i], headGrad, attnInput);
        WriteLn('        headGrad после attention: ', Length(headGrad), 'x', Length(headGrad[0]));
        
        // Накопление градиентов
        if (Length(headGrad) = Length(Result)) and (Length(headGrad[0]) = Length(Result[0])) then
        begin
          Result := MatrixAdd(Result, headGrad);
          WriteLn('        Градиенты головы добавлены');
        end
        else
        begin
          WriteLn('        Предупреждение: Несовпадение размеров headGrad');
          WriteLn('        headGrad: ', Length(headGrad), 'x', Length(headGrad[0]));
          WriteLn('        Result: ', Length(Result), 'x', Length(Result[0]));
          
          // Масштабируем до правильного размера
          headGrad := ScaleMatrixToSize(headGrad, Length(Result), Length(Result[0]));
          Result := MatrixAdd(Result, headGrad);
        end;
      end
      else
      begin
        WriteLn('        Предупреждение: Wo пустой');
      end;
      
    except
      on E: Exception do
      begin
        WriteLn('      ОШИБКА в голове ', i, ': ', E.Message);
        // Пропускаем эту голову
      end;
    end;
  end;
end;

// MatrixOps.pas - добавить эту функцию если её нет

function ScaleMatrixToSize(const m: TDoubleMatrix; newRows, newCols: Integer): TDoubleMatrix;
var
  i, j: Integer;
  rowScale, colScale: Double;
begin
  if (Length(m) = 0) or (newRows <= 0) or (newCols <= 0) then
  begin
    SetLength(Result, newRows, newCols);
    FillMatrix(Result, 0.0);
    Exit;
  end;
  
  SetLength(Result, newRows, newCols);
  rowScale := Length(m) / newRows;
  colScale := Length(m[0]) / newCols;
  
  for i := 0 to newRows - 1 do
  begin
    for j := 0 to newCols - 1 do
    begin
      if (Trunc(i * rowScale) < Length(m)) and 
         (Trunc(j * colScale) < Length(m[Trunc(i * rowScale)])) then
        Result[i][j] := m[Trunc(i * rowScale)][Trunc(j * colScale)]
      else
        Result[i][j] := 0.0;
    end;
  end;
end;